package com.leilei.source.customize;


import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
import org.apache.flink.api.common.RuntimeExecutionMode;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.SourceFunction;

import java.util.Random;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author lei
 * @version 1.0
 * @date 2021/3/8 22:33
 * @desc flink 自定义数据源
 */
public class FlinkSourceCustomize {
    public static void main(String[] args) throws Exception {
        //准备环境 env
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        env.setRuntimeMode(RuntimeExecutionMode.AUTOMATIC);
        //获取数据源 source
        DataStreamSource<VehicleAlarm> dataStreamSource = env.addSource(new VehicleAlarmSource());
        //数据处理
        SingleOutputStreamOperator<VehicleAlarm> streamOperator = dataStreamSource
                .map(alarm -> {
                    String[] colorArr = new String[]{"黄", "蓝"};
                    alarm.setColor(colorArr[(int) (Math.random() * 2)]);
                    return alarm;
                }).returns(VehicleAlarm.class);
        //数据收集 sink
        streamOperator.print();
        //程序执行 execute
        env.execute();

    }

    public static class VehicleAlarmSource implements SourceFunction<VehicleAlarm> {
        private Boolean flag = true;
        AtomicInteger idWork = new AtomicInteger(9999);

        @Override
        public void run(SourceContext<VehicleAlarm> ctx) throws Exception {
            while (flag) {
                Thread.sleep(5000);
                VehicleAlarm alarm = new VehicleAlarm(idWork.getAndDecrement(),
                        "川A000" + new Random().nextInt(10), "黄", UUID.randomUUID().toString());
                ctx.collect(alarm);
                System.out.println("加载数据....并休眠五秒");
            }
        }

        @Override
        public void cancel() {
            flag = false;
        }
    }

    @Data
    @NoArgsConstructor
    @AllArgsConstructor
    public static class VehicleAlarm {
        private Integer id;
        private String plate;
        private String color;
        private String alarmId;
    }
}
